# -*- coding: utf-8 -*-
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# Modified by Ke Sun (sunk@mail.ustc.edu.cn), Tianheng Cheng(tianhengcheng@gmail.com)
# ------------------------------------------------------------------------------

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import codecs
import csv
import os
import logging
import time
from pathlib import Path

import torch
import torch.optim as optim
import torchvision

import numpy as np
import cv2


def save_csv(predictions, filenames, save_path):
    '''

    :param predictions: n points of [h,w]
    :param filenames:
    :param save_path:
    :return:
    '''
    predictions = predictions.cpu()
    with codecs.open(save_path, 'w', 'utf-8') as csv_file:
        fieldnames = ['file_name', 'x0', 'x2', 'y0', 'y2']  # 表头
        writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
        writer.writeheader()
        for i, (prediction, filename) in enumerate(zip(predictions, filenames)):
            pts = dict()
            pts['file_name'] = filename
            pts['x0'] = prediction[0, 0].item()
            pts['x2'] = prediction[1, 0].item()
            pts['y0'] = prediction[0, 1].item()
            pts['y2'] = prediction[1, 1].item()
            writer.writerow(pts)

    csv_file.close()


def create_logger(cfg, cfg_name, phase='train'):
    root_output_dir = Path(cfg.OUTPUT_DIR)
    # set up logger
    if not root_output_dir.exists():
        print('=> creating {}'.format(root_output_dir))
        root_output_dir.mkdir()

    dataset = cfg.DATASET.DATASET
    model = cfg.MODEL.NAME
    cfg_name = os.path.basename(cfg_name).split('.')[0]

    final_output_dir = root_output_dir / dataset / cfg_name

    print('=> creating {}'.format(final_output_dir))
    final_output_dir.mkdir(parents=True, exist_ok=True)

    time_str = time.strftime('%Y-%m-%d-%H-%M')
    log_file = '{}_{}_{}.log'.format(cfg_name, time_str, phase)
    final_log_file = final_output_dir / log_file
    head = '%(asctime)-15s %(message)s'
    logging.basicConfig(filename=str(final_log_file),
                        format=head)
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    console = logging.StreamHandler()
    logging.getLogger('').addHandler(console)

    tensorboard_log_dir = Path(cfg.LOG_DIR) / dataset / model / \
                          (cfg_name + '_' + time_str)
    print('=> creating {}'.format(tensorboard_log_dir))
    tensorboard_log_dir.mkdir(parents=True, exist_ok=True)

    return logger, str(final_output_dir), str(tensorboard_log_dir)


def get_optimizer(cfg, model):
    optimizer = None
    if cfg.TRAIN.OPTIMIZER == 'sgd':
        optimizer = optim.SGD(
            filter(lambda p: p.requires_grad, model.parameters()),
            lr=cfg.TRAIN.LR,
            momentum=cfg.TRAIN.MOMENTUM,
            weight_decay=cfg.TRAIN.WD,
            nesterov=cfg.TRAIN.NESTEROV
        )
    elif cfg.TRAIN.OPTIMIZER == 'adam':
        optimizer = optim.Adam(
            filter(lambda p: p.requires_grad, model.parameters()),
            lr=cfg.TRAIN.LR
        )
    elif cfg.TRAIN.OPTIMIZER == 'rmsprop':
        optimizer = optim.RMSprop(
            filter(lambda p: p.requires_grad, model.parameters()),
            lr=cfg.TRAIN.LR,
            momentum=cfg.TRAIN.MOMENTUM,
            weight_decay=cfg.TRAIN.WD,
            alpha=cfg.TRAIN.RMSPROP_ALPHA,
            centered=cfg.TRAIN.RMSPROP_CENTERED
        )

    return optimizer


def save_checkpoint(states, predictions, is_best,
                    output_dir, filename='checkpoint.pth'):
    preds = predictions.cpu().data.numpy()
    torch.save(states, os.path.join(output_dir, filename))
    torch.save(preds, os.path.join(output_dir, 'current_pred.pth'))

    latest_path = os.path.join(output_dir, 'latest.pth')
    if os.path.islink(latest_path):
        os.remove(latest_path)
    # os.symlink(os.path.join(output_dir, filename), latest_path)
    os.symlink(filename, latest_path)

    if is_best and 'state_dict' in states.keys():
        torch.save(states['state_dict'].module, os.path.join(output_dir, 'model_best.pth'))


def cat_score_map(score_map):
    score_imgs = []
    batch = score_map.size(0)

    for i in range(batch):
        score_imgs.append(score_map[i, 0, :, :].unsqueeze(0))
        score_imgs.append(score_map[i, 1, :, :].unsqueeze(0))
    imgs = torchvision.utils.make_grid(score_imgs, nrow=2, padding=0)
    return imgs


def log_imgs(imgs, writer_dict):
    writer = writer_dict['writer']
    writer.add_image('score_map', imgs)


def get_CAM(img, heatmap):
    heatmap = heatmap.cpu().numpy()
    heatmap = np.mean(heatmap, axis=0)

    heatmap = np.maximum(heatmap, 0)
    heatmap /= np.max(heatmap)

    heatmap = cv2.resize(heatmap, (img.size(1), img.size(2)))  # 将热力图的大小调整为与原始图像相同
    heatmap = np.uint8(255 * heatmap)  # 将热力图转换为RGB格式
    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)

    heatmap = np.transpose(heatmap, [2, 1, 0])
    superimposed_img = torch.tensor(heatmap) * 0.8 + img  # 将热力图应用于原始图像,这里的0.4是热力图强度因子
    return superimposed_img


def log_tpts_imgs(inp, meta, preds, mean, std, writer_dict):
    '''
    :param inp: transformed images, [b,c,w,h]
    :return:
    '''
    batch_size = inp.size(0)
    inp = inp.cpu().numpy()
    pts = meta['pts'].cpu().numpy()
    preds = preds.cpu().numpy()
    imgs = []
    gt_color = (0, 0, 255)  # blue
    predict_color = (255, 0, 0)  # red
    for i in range(batch_size):
        img = np.transpose(inp[i, :, :, :], [1, 2, 0])  # w,h,c
        img *= std
        img += mean
        img *= 255
        img = img.astype(np.float).copy()
        cv2.circle(img, center=(pts[i, 0, 0], pts[i, 0, 1]), radius=5, color=gt_color)
        cv2.circle(img, center=(pts[i, 1, 0], pts[i, 1, 1]), radius=10, color=gt_color)
        cv2.circle(img, center=(preds[i, 0, 0], preds[i, 0, 1]), radius=5, color=predict_color)
        cv2.circle(img, center=(preds[i, 1, 0], preds[i, 1, 1]), radius=10, color=predict_color)
        img = np.transpose(img, [2, 1, 0])  # c,h,w
        imgs.append(img)
    imgs = np.stack(imgs)
    imgs = torchvision.utils.make_grid(torch.tensor(imgs), nrow=4, padding=5, normalize=True)
    log_imgs(imgs, writer_dict)
